#include "dyn_roots_multiDim_staPlus.h"

STAPlus::STAPlus(MultiMaps &map,
                 colvec init,
                 const double tol,
                 const double tol2,
                 const double para,
                 const int period,
                 const int Nmax)
    :MultiSTA(map, init, tol, para, period,Nmax)
{
    setTolerance2(tol2);
    setPeriod(period);
}

int STAPlus::solve()
{
    // We start the clock
    start = clock();

    // We initialize the return value.
    int returnCode(1);
    int orbitFound(0);

    int dim = m_initialValue.n_elem;
    fixedPoints.set_size(period,dim);
    rowvec transient;

    transient = m_initialValue.t();

    // Actual algorithm
    for (unsigned int i=0;i<reflectionMatrices.n_slices;i++)
    {
        for (int j=0;j<maxIterations;j++)
        {
            rowvec dx = computeCorrection(transient.t(),
                                       reflectionMatrices.slice(i),
                                       lambda,
                                       period,
                                          dim).t();

            transient += dx;

            if (norm(dx,2) < tolerance)
            {
                fixedPoints.row(0) = transient;
                orbitFound = 1;
                break;
            }
        }

        if (orbitFound)
        {
            for (int k=1;k<period;k++)
            {
                fixedPoints.row(k) = multiMap.iterate(fixedPoints.row(k-1).t(),1).t();

//                for (int j=0;j<maxIterations;j++)
//                {
//                    rowvec dx = computeCorrection(transient.t(),reflectionMatrices.slice(i),0,period,dim).t();
//                    transient += dx;;

//                    if (norm(dx,2) < tolerance2)
//                    {
//                        fixedPoints.row(k) = transient;
//                        break;
//                    }
//                }
            }
            returnCode = 0;
            break;
        }
    }
    return returnCode;
}

colvec STAPlus::computeCorrection(colvec init, mat C, double para, int p, int dim)
{
    colvec iterationValue = (multiMap.iterate(init,p)-init);
    mat invertedMatrix = (para*norm(iterationValue,2)*eye(dim,dim)-C*multiMap.jacobian(init,p));
    colvec dx = arma::solve(invertedMatrix,C*iterationValue);

    return dx;
}
